#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
tools/cosy_infer.py — CosyVoice2.inference_sft 直连 + speaker 自适配
- 优先用用户传入的 --spk
- 若失败，自动从 {model_dir}/cosyvoice2.yaml 中读取可用 speaker 列表并逐一尝试
- 全程打印详细调试信息到 stderr，便于在 pipeline 外单独运行排查
"""
from __future__ import annotations
import argparse, os, sys, io, wave, contextlib, importlib, traceback
import numpy as np

def _ensure_sys_path():
    src = os.environ.get("COSYVOICE_SRC") or "/root/AI_boardcast/CosyVoice"
    if src and os.path.isdir(src) and src not in sys.path:
        sys.path.insert(0, src)

def _to_wav_bytes(samples, sr: int) -> bytes:
    if isinstance(samples, (bytes, bytearray)):
        return bytes(samples)
    if isinstance(samples, tuple) and len(samples) == 2:
        samples, sr = samples
    arr = np.asarray(samples, dtype=np.float32)
    arr = np.clip(arr, -1.0, 1.0)
    buf = io.BytesIO()
    with contextlib.closing(wave.open(buf, "wb")) as wf:
        wf.setnchannels(1); wf.setsampwidth(2); wf.setframerate(int(sr))
        wf.writeframes((arr * 32767.0).astype("<i2").tobytes())
    return buf.getvalue()

def _load_yaml_speakers(model_dir: str) -> list[str]:
    """从 cosyvoice2.yaml 中尽可能挖 speaker 候选名"""
    import yaml
    yaml_path = os.path.join(model_dir, "cosyvoice2.yaml")
    cands = []
    try:
        if os.path.isfile(yaml_path):
            data = yaml.safe_load(open(yaml_path, "r", encoding="utf-8"))
            # 常见位置 1：spk、speaker、spk_list
            for k in ("spk", "speaker", "speakers", "spk_list", "speaker_list"):
                v = data.get(k)
                if isinstance(v, list):
                    cands += [str(x) for x in v]
                elif isinstance(v, dict):
                    cands += [str(x) for x in v.keys()]
            # 常见位置 2：nested
            for k in ("tts", "sft", "voice", "inference"):
                if isinstance(data.get(k), dict):
                    for kk in ("spk", "speakers", "spk_list", "speaker_list"):
                        v = data[k].get(kk)
                        if isinstance(v, list):
                            cands += [str(x) for x in v]
                        elif isinstance(v, dict):
                            cands += [str(x) for x in v.keys()]
    except Exception as e:
        print(f"[debug] parse yaml failed: {repr(e)}", file=sys.stderr)
    # 去重，保留顺序
    seen = set(); out = []
    for x in cands:
        if x and x not in seen:
            seen.add(x); out.append(x)
    return out

def main():
    _ensure_sys_path()
    ap = argparse.ArgumentParser()
    ap.add_argument("--model_dir", required=True)
    ap.add_argument("--spk", default=None)
    ap.add_argument("--text", required=True)
    ap.add_argument("--wav", required=True)
    ap.add_argument("--sr", type=int, default=24000)
    args = ap.parse_args()

    # 1) 导入模块 & 类
    try:
        cv_mod = importlib.import_module("cosyvoice.cli.cosyvoice")
        print(f"[info] cosyvoice module: {getattr(cv_mod, '__file__', None)}", file=sys.stderr)
    except Exception as e:
        print(f"[fatal] import cosyvoice.cli.cosyvoice failed: {repr(e)}", file=sys.stderr)
        sys.exit(3)

    CV = getattr(cv_mod, "CosyVoice2", None) or getattr(cv_mod, "CosyVoice", None)
    if CV is None:
        print("[fatal] no CosyVoice2/CosyVoice class", file=sys.stderr); sys.exit(3)

    # 2) 实例化
    try:
        tts = CV(model_dir=args.model_dir)
        print(f"[info] instantiated {CV.__name__} with model_dir={args.model_dir}", file=sys.stderr)
    except Exception as e:
        print(f"[fatal] instantiate {CV.__name__} failed: {repr(e)}", file=sys.stderr)
        traceback.print_exc()
        sys.exit(3)

    # 3) 拿到 inference_sft
    infer = getattr(tts, "inference_sft", None)
    if not callable(infer):
        print("[fatal] CosyVoice2 has no inference_sft method", file=sys.stderr)
        sys.exit(3)

    # 4) 组装 spk 候选：用户输入 → 常见占位 → YAML 解析
    spk_candidates = []
    if args.spk: spk_candidates.append(args.spk)
    spk_candidates += ["female_0001", "male_0001", "cn_female_1", "cn_male_1"]
    spk_candidates += _load_yaml_speakers(args.model_dir)
    # 去重
    seen = set(); spk_candidates = [s for s in spk_candidates if not (s in seen or seen.add(s))]
    if not spk_candidates:
        spk_candidates = [None]  # 少数模型支持省略 spk_id

    last_err = None
    for spk_try in spk_candidates:
        try:
            print(f"[info] try spk_id={spk_try!r}", file=sys.stderr)
            audio = infer(tts_text=args.text, spk_id=spk_try, stream=False, speed=1.0, text_frontend=True)
            wav_bytes = _to_wav_bytes(audio, args.sr)
            os.makedirs(os.path.dirname(args.wav), exist_ok=True)
            with open(args.wav, "wb") as f:
                f.write(wav_bytes)
            print(f"[ok] -> {args.wav} (spk_id={spk_try!r})", file=sys.stderr)
            return
        except Exception as e:
            last_err = e
            print(f"[warn] inference_sft failed with spk_id={spk_try!r}: {repr(e)}", file=sys.stderr)
            continue

    print(f"[fatal] all spk candidates failed. last_err={repr(last_err)}", file=sys.stderr)
    sys.exit(3)

if __name__ == "__main__":
    main()
